"""
数据库模块 - 声明式事务管理

此模块提供了 `@transactional` 装饰器，旨在为业务逻辑方法提供
类似于 Spring 的声明式事务能力。
"""
from __future__ import annotations
import inspect
from contextvars import ContextVar
from functools import wraps
from typing import Any, Callable, Optional, Type, Tuple
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, AsyncTransaction

from axiom_boot.di import autowired

# ContextVar 用于在异步任务的上下文中安全地传递会话。
# 这是实现“无感知”会话传播的关键。
_current_session: ContextVar[Optional[AsyncSession]] = ContextVar("_current_session", default=None)


class TransactionManager:
    """
    事务管理器，负责处理单个事务的完整生命周期。
    它作为一个异步上下文管理器使用，并支持嵌套事务 (SAVEPOINT)。
    """
    def __init__(
        self,
        session_factory: async_sessionmaker[AsyncSession],
        read_only: bool = False,
        isolation_level: Optional[str] = None,
        rollback_for: Optional[Tuple[Type[Exception], ...]] = None,
        no_rollback_for: Optional[Tuple[Type[Exception], ...]] = None
    ):
        if not callable(session_factory):
            raise TypeError("session_factory 必须是一个可调用的对象。")
        self.session_factory = session_factory
        self.read_only = read_only
        self.isolation_level = isolation_level
        self.rollback_for = rollback_for or ()
        self.no_rollback_for = no_rollback_for or ()
        self._session_token = None
        self._transaction: Optional[AsyncTransaction] = None
        self._is_nested = False

    async def __aenter__(self) -> AsyncSession:
        """进入上下文时，获取会话并开启事务或子事务 (SAVEPOINT)。"""
        existing_session = _current_session.get()

        if existing_session and existing_session.in_transaction():
            # 如果已存在事务，则创建嵌套事务 (SAVEPOINT)
            self._is_nested = True
            self.session = existing_session
            self._transaction = self.session.begin_nested()
        else:
            # 否则，创建新的会话和顶级事务
            self.session = self.session_factory()
            self._session_token = _current_session.set(self.session)
            self._transaction = self.session.begin()
            
            execution_options = {}
            if self.isolation_level:
                execution_options["isolation_level"] = self.isolation_level
            if self.read_only:
                execution_options["postgresql_readonly"] = True
            
            if execution_options:
                await self.session.connection(execution_options=execution_options)

        await self._transaction.__aenter__()
        return self.session

    async def __aexit__(self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[Any]):
        """退出上下文时，根据结果和回滚规则提交或回滚事务/子事务。"""
        if self._transaction is None:
            return

        should_rollback = exc_type is not None
        if should_rollback:
            # 检查 no_rollback_for 规则
            if isinstance(exc_val, self.no_rollback_for):
                should_rollback = False
            # 检查 rollback_for 规则 (如果 no_rollback_for 未匹配)
            elif self.rollback_for and isinstance(exc_val, self.rollback_for):
                should_rollback = True
        
        final_exc_type = exc_type if should_rollback else None
        final_exc_val = exc_val if should_rollback else None
        final_exc_tb = exc_tb if should_rollback else None

        try:
            # 将最终的异常状态传递给事务上下文管理器进行处理
            await self._transaction.__aexit__(final_exc_type, final_exc_val, final_exc_tb)
        finally:
            # 如果是顶级事务，则在退出后关闭会话并重置 ContextVar
            if not self._is_nested and self._session_token:
                await self.session.close()
                _current_session.reset(self._session_token)


def transactional(
    func: Optional[Callable] = None,
    *,
    read_only: bool = False,
    isolation_level: Optional[str] = None,
    rollback_for: Optional[Tuple[Type[Exception], ...]] = None,
    no_rollback_for: Optional[Tuple[Type[Exception], ...]] = None,
):
    """
    一个方法装饰器，为被装饰的异步方法提供声明式事务管理。
    """
    if func is None:
        return lambda f: transactional(
            f, 
            read_only=read_only, 
            isolation_level=isolation_level,
            rollback_for=rollback_for,
            no_rollback_for=no_rollback_for
        )

    @wraps(func)
    async def wrapper(*args, **kwargs):
        service_instance = args[0]
        session_factory = getattr(service_instance, 'db_session_factory', None)

        if not session_factory:
            raise AttributeError(
                f"无法在服务 '{service_instance.__class__.__name__}' 上找到 'db_session_factory' 属性。"
                "请确保 DatabaseProcessor 已正确运行并注入了会话工厂。"
            )

        manager = TransactionManager(
            session_factory=session_factory,
            read_only=read_only,
            isolation_level=isolation_level,
            rollback_for=rollback_for,
            no_rollback_for=no_rollback_for,
        )
        
        # 确保只对异步函数生效
        if not inspect.iscoroutinefunction(func):
            raise TypeError("`@transactional` 装饰器只能用于异步 (async def) 方法。")
            
        async with manager as session:
            # 这里不再需要手动注入 session，因为 BaseMapper 会通过 _current_session.get() 获取它
            return await func(*args, **kwargs)
            
    return wrapper 